# This convient library came from https://github.com/BetaRavener/upy-websocket-server
# Thanks very much for this lab and comprehensive demo

import os, time
from machine import UART, Pin
import gc

from ws_connection import ClientClosedError
from ws_multiserver import WebSocketMultiServer
from ws_server import WebSocketClient

from ws_app_motor import WsModbusClient



class WsUartClient(WebSocketClient):
    """实现串口透传功能(Pass through data between UART and Websocket)."""
    def __init__(self, conn):
        super().__init__(conn)
        # even parity(parity=0)
        self.uart_isp = UART(1, baudrate=115200, parity=0, stop=1, rx=25, tx=26, rxbuf=256)
    
    def process(self):
        # ===== Send received data to UART0 =====
        try:
            _recv = self.connection.read()
        except ClientClosedError:
            self.connection.close()
            return
        
        if _recv is not None:
            print("isp tx>|%s"%_recv)
            self.uart_isp.write(_recv)
        
        # ===== Send UART0 data to Websocket =====
        if self.uart_isp.any():
            _resp = self.uart_isp.read()
            _resp = ''.join(['%02X'%x for x in _resp])  # convert to hex
            print("isp rx<|%s"%_resp)
            self.connection.write(_resp)


class WsCommandClient(WebSocketClient):
    """Perform some command from websocket"""
    def __init__(self, conn):
        super().__init__(conn)
        self._ACK = b'0x79'
        self._NACK = b'0x78'
        self._cmd_map = (
            self._do_gpio_output,  # 0
        )
    
    def _do_gpio_output(self, data):
        try:
            _gpio_num = data[0]
            _gpio_level = data[1]
        except IndexError:
            return  # stop execution
        
        _p = Pin(_gpio_num, Pin.OUT)
        _p.value(_gpio_level)
        del _p
        
        self.connection.write(self._ACK)
    
    def process(self):
        try:
            _recv = self.connection.read()
        except ClientClosedError:
            self.connection.close()
            return
        
        if _recv is not None:
            try:
                _cmd_idx, _data = _recv[0], _recv[1:]
                self._cmd_map[_cmd_idx](_data)
            except IndexError:
                pass


class WsDoNothingClient(WebSocketClient):
    """A placeholder to delay websocket connection close, so the client
    could be cleaned togather."""
    def __init__(self, conn):
        super().__init__(conn)
    
    def process(self):
        self.connection.close()


class WsClientFactory:
    """
    '0x10 01' -> Uart Client
    '0x10 02' -> Command Client
    """
    def __init__(self, conn):
        self._HEAD = 0x10
        self._TYPE_UART = 0x01
        self._TYPE_COMMAND = 0x02
        self._TYPE_MODBUS = 0x03
        self._TIMEOUT = 3  # second

        self.connection = conn
    
    def make_client(self):
        _time_start = time.time()
        _recv = bytearray()
        _idx = 0
        while (time.time() - _time_start) < self._TIMEOUT:
            # collect data
            _dat = self.connection.read()
            if _dat is None:
                continue
            
            _recv.extend(_dat)
            if len(_recv) < 2:
                continue
            
            # process data
            if _recv[_idx] == self._HEAD:
                _idx += 1
                if _recv[_idx] == self._TYPE_COMMAND:
                    print("command")  # debug
                    return WsCommandClient(self.connection)
                elif _recv[_idx] == self._TYPE_UART:
                    print("uart")   # debug
                    return WsUartClient(self.connection)
                elif _recv[_idx] == self._TYPE_MODBUS:
                    print("modbus")
                    return WsModbusClient(self.connection)
        
        # no expected request until time out
        print("do nothing", _recv)  # debug
        return WsDoNothingClient(self.connection)
                

class IspWsServer(WebSocketMultiServer):
    def __init__(self):
        super().__init__('index.html', 3)

    def _make_client(self, conn):   
        print("make conn")
        _factory = WsClientFactory(conn) 

        return _factory.make_client()


status_pin = Pin(2, Pin.OUT)
last_t = time.time()
def status_blink():
    """Blink to indicate ESP32 is working..."""
    global last_t
    if time.time() - last_t > 0.5:
        status_pin.value(not status_pin.value())
        last_t = time.time()  # negelect overflow


server = IspWsServer()
server.start()

try:
    while True:
        server.process_all()
        status_blink()
        gc.collect()
except KeyboardInterrupt:
    print("KeyboardInterrupt")

server.stop()